package udp

import (
	"context"
	"encoding/binary"
	"fmt"
	"github.com/kataras/golog"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"net"
	"nprobe/core"
	"nprobe/utils/routines"
	"time"
)

type UDPReceiver struct {
	core.BaseReceiver
	Bucket *core.PktRecordBucket
}

func NewUDPReceiver(
	ctx context.Context, name string, id int, task *core.ProbeTask, salts map[int][]byte, targetSet map[string]bool, bucket *core.PktRecordBucket,
) *UDPReceiver {
	return &UDPReceiver{
		BaseReceiver: core.BaseReceiver{
			Ctx:      ctx,
			Name:     name,
			ID:       id,
			Task:     task,
			RBuf:     nil,
			Salts:    salts,
			DstIPSet: targetSet,
		},
		Bucket: bucket,
	}
}

func (r *UDPReceiver) StartReceiver() error {
	routines.GoWithRecover(r.Name, r.doReceiveICMP)
	return nil
}

func (r *UDPReceiver) doReceiveICMP() {
	conn, err := core.OpenReceiverIP4PacketConn("ip4:icmp", "0.0.0.0") // 本地地址
	if err != nil {
		return
	}
	r.Conn = conn

	filter := core.CreateBPFUDPICMP(r.Task.SrcPortMin, r.Task.SrcPortMax, r.Task.DstPortMin, r.Task.DstPortMax)
	err = core.SetRawConnBPF(conn, filter)
	if err != nil {
		fmt.Println("unable to open network connection")
		return
	}

	payloadBuf := core.AllocRxBuf(r.Task.MsgLen)
	for {
		select {
		case <-r.Ctx.Done():
			conn.Close()
			return
		default:
			r.receiveOneICMPPkt(payloadBuf)
		}
	}
}

func (r *UDPReceiver) receiveOneICMPPkt(payloadBuf []byte) {
	conn := r.Conn
	err := conn.SetReadDeadline(time.Now().Add(time.Second))
	if err != nil {
		return
	}

	_, payload, _, err := conn.ReadFrom(payloadBuf)
	if err != nil {
		golog.Debugf("read packet failed. err: %+v", err)
		return
	}
	rxTs := time.Now().UnixNano()

	icmpMsg, err := icmp.ParseMessage(1, payload)

	var originIPPkt []byte
	switch icmpMsg.Type {
	case ipv4.ICMPTypeDestinationUnreachable:
		originIPPkt = icmpMsg.Body.(*icmp.DstUnreach).Data
		r.processICMPEchoReply(rxTs, originIPPkt)
	}
}

func (r *UDPReceiver) processICMPEchoReply(rxTs int64, payload []byte) {

	ipVer := (payload[0] >> 4) & 0x0f
	if ipVer != 4 {
		return
	}

	hLen := (payload[0] & 0x0f) * 4
	srcIP := net.IPv4(payload[12], payload[13], payload[14], payload[15])
	dstIP := net.IPv4(payload[16], payload[17], payload[18], payload[19])
	ipID := binary.BigEndian.Uint16(payload[4:6])
	srcPort := binary.BigEndian.Uint16(payload[hLen : hLen+2])
	dstPort := binary.BigEndian.Uint16(payload[hLen+2 : hLen+4])

	if !r.DstIPSet[dstIP.String()] {
		return
	}

	if srcPort > r.Task.SrcPortMax || srcPort < r.Task.SrcPortMin || dstPort > r.Task.DstPortMax || dstPort < r.Task.DstPortMin {
		return
	}

	r.recordPkt(srcIP.To4().String(), dstIP.To4().String(), srcPort, dstPort, ipID, 0, rxTs, false)
}

func (r *UDPReceiver) recordPkt(localIP string, remoteIP string, srcPort uint16, dstPort uint16, ipID uint16, txTs int64, rxTs int64, bitFlip bool) {

	key := fmt.Sprintf("%s-%d-%d-%d", remoteIP, ipID, srcPort, dstPort)
	golog.Debugf("%s: %s:%d->%s:%d, ipID: %d, txTs: %d", r.Name, localIP, srcPort, remoteIP, dstPort, ipID, rxTs)
	pktRecord := &core.PktRecord{
		Key:      key,
		SrcIP:    localIP,
		DstIP:    remoteIP,
		Proto:    r.Task.Proto,
		Tos:      r.Task.TOS,
		RxTxTs:   txTs,
		RxTs:     rxTs,
		Received: true,
		BitFlip:  bitFlip,
	}

	r.Bucket.Update(pktRecord)
}
